﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.Linq;
using System.Text;
using System.Windows.Forms;
using GMap.NET.WindowsForms;
using GMap.NET;
using GMap.NET.MapProviders;
using EkspSys.Network;
using System.Net;
using EkspSys.network;
using EkspSys.Map;
using GMap.NET.WindowsForms.Markers;
using System.Diagnostics;
using System.Collections.Concurrent;
using System.Threading;
using EkspSys.Components;

namespace EkspSys
{
    public partial class MainWindow : Form
    {
        private WinPcap winPcap = new WinPcap();
        private BackgroundWorker worker = new BackgroundWorker();

        BlockingCollection<Package> packageQueue;

        // TODO: Refactor out in its own class
        public class ProcessComparer : IEqualityComparer<Process>
        {
            public bool Equals(Process x, Process y)
            {
                return x.Id == y.Id;
            }

            public int GetHashCode(Process p)
            {
                return p.Id;
            }
        }

        private Dictionary<GMapMarker, EndPointInfo> markerEndPointMap = new Dictionary<GMapMarker,EndPointInfo>();
        private Dictionary<IPAddress, GMapMarker> ipToMarkerMap = new Dictionary<IPAddress, GMapMarker>();
        private Dictionary<Process, GMapOverlay> processToOverlayMap = new Dictionary<Process, GMapOverlay>(new ProcessComparer());

        private Dictionary<IPAddress, Tuple<GMapOverlay, IListItem>> ipOverlayTraceRTMap
            = new Dictionary<IPAddress, Tuple<GMapOverlay, IListItem>>();

        private GMapOverlay unknownOverlay;

        public MainWindow()
        {
            InitializeComponent();
            InitializePackageSniffing();
            InitializeMap();
            InitializeAppList();
            InitlalizeNetworkInterfaceChanger();
        }

        private void InitializePackageSniffing()
        {
            worker.WorkerSupportsCancellation = true;
            worker.WorkerReportsProgress = true;
            worker.ProgressChanged += new ProgressChangedEventHandler(WorkerChanges);
            worker.DoWork += new DoWorkEventHandler(DoHandlePackage);

            winPcap.PackageSniffed += new PackageSniffedEventHandler(PackageSniffedHandler);
        }

        private void InitializeMap()
        {
            mapexplr.MapProvider = GMapProviders.BingHybridMap;
            mapexplr.MinZoom = 1;
            mapexplr.MaxZoom = 17;
            mapexplr.Zoom = 2;
            mapexplr.Manager.Mode = AccessMode.ServerAndCache;
            mapexplr.DragButton = MouseButtons.Left;

            mapexplr.OnMarkerClick += new MarkerClick(MarkerClicked);

            unknownOverlay = new GMapOverlay(mapexplr, "UnknownOverlay");
            
            mapexplr.Overlays.Add(unknownOverlay);
        }

        private void InitializeAppList()
        {
            StringAppListItem item = new StringAppListItem("Alle", new EkspSys.Components.StringAppListItem.ListItemAction(delegate() {
                tracesList.ClearSelected();
                SetEndPointOverlayVisibility(true);
                HideTraceRouteOverlays();
            }));

            appList.Items.Add(item);
            appList.SelectedItem = item;

            tracesList.ClearSelected();
        }

        private void SetEndPointOverlayVisibility(bool visible)
        {
            newEndPointOverlaysVisible = visible;
            unknownOverlay.IsVisibile = visible;
            foreach (GMapOverlay overlay in processToOverlayMap.Values)
                overlay.IsVisibile = visible;
        }

        private void HideTraceRouteOverlays()
        {
            foreach (Tuple<GMapOverlay, IListItem> val in ipOverlayTraceRTMap.Values)
                val.Item1.IsVisibile = false;
        }

        private void HideAllOverlays()
        {
            SetEndPointOverlayVisibility(false);
            HideTraceRouteOverlays();
        }

        private void InitlalizeNetworkInterfaceChanger()
        {
            List<NetworkInterface> interfaces = WinPcap.GetNetworkInterfaces();

            if (interfaces.Count != 0)
            {
                NetworkInterface defNif = interfaces[0];
                networkInterfaceChanger.Text = defNif.description;
                NetworkInterfaceChanged(defNif)(null, null);
            }

            foreach (NetworkInterface nif in interfaces) {
                // TODO: Add a nice image
                networkInterfaceChanger.DropDownItems.Add(nif.description, global::EkspSys.Properties.Resources.networkcard, NetworkInterfaceChanged(nif));
            }
        }

        private AutoResetEvent workerFinished = new AutoResetEvent(false);

        private CancellationTokenSource cancellationTokenSource;
        private CancellationToken cancellationToken;

        private EventHandler NetworkInterfaceChanged(NetworkInterface nif)
        {
            return delegate(object sender, EventArgs e) {
                // Stop package captureing and worker thread
                winPcap.StopCapture();
                
                if (worker.IsBusy)
                {
                    cancellationTokenSource.Cancel();
                    worker.CancelAsync();
                    workerFinished.WaitOne();
                    
                    // XXX: HACK! In need of more time...
                    // Do form events to prevent deadlock
                    this.Enabled = false;
                    while (worker.IsBusy)
                    {
                        Application.DoEvents();
                    }
                    this.Enabled = true;
                }

                // Add a new package queue for worker to work on
                packageQueue = new BlockingCollection<Package>();

                // Clean current markers from the map
                markerEndPointMap.Clear();
                ClearMapOverlays();

                // Start a new capture session
                cancellationTokenSource = new CancellationTokenSource();
                cancellationToken = cancellationTokenSource.Token;
                worker.RunWorkerAsync();
                winPcap.StartCaptureAsync(nif);

                networkInterfaceChanger.Text = nif.description;
            };
        }

        private void ClearMapOverlays()
        {
            unknownOverlay.Markers.Clear();
            foreach (GMapOverlay overlay in processToOverlayMap.Values)
            {
                overlay.Markers.Clear();
            }
        }

        private void PackageSniffedHandler(Package pck)
        {
            packageQueue.Add(pck);
        }

        private delegate void UpdateUIExecutor();

        private void WorkerChanges(object sender, ProgressChangedEventArgs e)
        {
            // Execute update in UI thread
            (e.UserState as UpdateUIExecutor)();
        }

        private bool newEndPointOverlaysVisible = true;

        private void DoHandlePackage(object sender, DoWorkEventArgs e)
        {
            while (!cancellationToken.IsCancellationRequested)
            {
                Package pck;
                try
                {
                    pck = packageQueue.Take(cancellationToken);
                }
                catch (OperationCanceledException _)
                {
                    // We have been canceled
                    break;
                }

                worker.ReportProgress(0, new UpdateUIExecutor(delegate() {
                    queuedPackages.Text = "Queued packages: " + packageQueue.Count();
                }));

                Tuple<IPAddress, ushort, PckDirection> remote = Utils.Get().GetRemoteSide(pck);
                if (remote == null)
                    continue;

                IPAddress addr = remote.Item1;
                ushort port = remote.Item2;

                Location location = GeoLookup.Get().LookupIPLocation(addr);
                if (location != null)
                {
                    Process process = null;
                    GMapMarker marker;

                    if (ipToMarkerMap.ContainsKey(addr)) {
                        marker = ipToMarkerMap[addr];
                    } else {
                        Tuple<IPAddress, ushort, PckDirection> local = Utils.Get().GetLocalSide(pck);
                        process = ProcessHelper.GetProcessForPort(pck.protocol, local.Item2);

                        marker = new GMapEndPointMarker(new PointLatLng(location.latitude, location.longitude),
                            ColorPool.Get().GetColor(process));
                        ipToMarkerMap.Add(addr, marker);

                        if (process == null)
                        {
                            if (unknownOverlay.Markers.Count == 0)
                            {
                                worker.ReportProgress(0, new UpdateUIExecutor(delegate() {
                                    appList.Items.Add(new StringAppListItem("Ukendte programmer",
                                        new StringAppListItem.ListItemAction(delegate() {
                                            tracesList.ClearSelected();
                                            HideAllOverlays();
                                            unknownOverlay.IsVisibile = true;
                                        })));
                                }));
                            }

                            // We could not find the process, add to unknown overlay
                            unknownOverlay.Markers.Add(marker);
                        }
                        else
                        {
                            GMapOverlay overlay;
                            if (processToOverlayMap.ContainsKey(process))
                            {
                                overlay = processToOverlayMap[process];
                            }
                            else
                            {
                                if (process.Id != 0)
                                {
                                    process.EnableRaisingEvents = true;
                                    process.Exited += new EventHandler(CleanupProcess);
                                }

                                overlay = new GMapOverlay(mapexplr, "Overlay" + process.Id);
                                overlay.IsVisibile = newEndPointOverlaysVisible;
                                processToOverlayMap.Add(process, overlay);

                                mapexplr.Overlays.Add(overlay);

                                // Update the UI in the UI thread
                                worker.ReportProgress(0, new UpdateUIExecutor(delegate() {
                                    int pos = appList.Items.Count;
                                    if (unknownOverlay.Markers.Count != 0)
                                        pos--;

                                    appList.Items.Insert(pos, new ApplicationListItem(process, ColorPool.Get().GetColor(process)));
                                } ));
                            }

                            if (!overlay.Markers.Contains(marker))
                                overlay.Markers.Add(marker);
                        }
                    }

                    EndPointInfo endPointInfo;
                    if (markerEndPointMap.ContainsKey(marker))
                    {
                        endPointInfo = markerEndPointMap[marker];
                    }
                    else
                    {
                        endPointInfo = new EndPointInfo(addr);
                        markerEndPointMap.Add(marker, endPointInfo);
                    }

                    ProcessInfo procInfo = endPointInfo.GetProcessInfo(port);
                    if (procInfo == null)
                    {
                        procInfo = new ProcessInfo();
                        procInfo.Process = process;
                        procInfo.Protocol = pck.protocol;
                        endPointInfo.SetProcessInfo(port, procInfo);
                    }
                    procInfo.Packets++;
                    procInfo.Data += pck.len;
                }
                else
                {
                    // We could not locate the IP
                    // Console.WriteLine("Unknown location for {0}", remote.Item1);
                }
            }

            workerFinished.Set();
        }

        private void CleanupProcess(object sender, EventArgs e)
        {
            Process proc = sender as Process;

            int pos = -1;
            for (int i = 0; i < appList.Items.Count; i++)
            {
                if (appList.Items[i] is ApplicationListItem)
                {
                    ApplicationListItem item = appList.Items[i] as ApplicationListItem;
                    if (item.Process.Id == proc.Id)
                    {
                        pos = i;
                        break;
                    }
                }
            }

            if (pos != -1)
            {
                worker.ReportProgress(0, new UpdateUIExecutor(delegate()
                {
                    appList.Items.RemoveAt(pos);
                }));
            }

            if (processToOverlayMap.ContainsKey(proc))
            {
                GMapOverlay overlay = processToOverlayMap[proc];
                overlay.Markers.Clear();
                mapexplr.Overlays.Remove(overlay);

                processToOverlayMap.Remove(proc);
            }
        }

        private GMapSelectableMarker   CurrentlySelectedMarker      = null;
        private EndPointInfoForm       CurrentlyDisplayedEndPointMarkerForm = null;
        private TraceRouteNodeInfoForm CurrentlyDisplayedRouteMarkerForm = null;

        private void RedrawMap()
        {
            mapexplr.Invalidate();//TODO: optimize?
        }

        private void OpenMarkerInfo()
        {
            if (CurrentlySelectedMarker == null) return;

            // TODO: Fix this when more time is available!!!
            if (CurrentlySelectedMarker is GMapEndPointMarker)
            {
                CurrentlyDisplayedEndPointMarkerForm = new EndPointInfoForm(markerEndPointMap[CurrentlySelectedMarker]);
                CurrentlyDisplayedEndPointMarkerForm.TraceRouteEvent += new EndPointInfoForm.TraceRouteDelegate(DoTraceRoute);
                CurrentlyDisplayedEndPointMarkerForm.FormClosed += new FormClosedEventHandler(delegate(object sender, FormClosedEventArgs fce)
                {
                    ClearMarkerInfo();
                    SelectMarker(null);
                    OpenMarkerInfo();
                });
                CurrentlyDisplayedEndPointMarkerForm.Show();
            }
            else if (CurrentlySelectedMarker is GMapRouteMarker)
            {
                GMapRouteMarker routeMarker = CurrentlySelectedMarker as GMapRouteMarker;
                CurrentlyDisplayedRouteMarkerForm = new TraceRouteNodeInfoForm(routeMarker.Address);
                CurrentlyDisplayedRouteMarkerForm.Show();
            }
        }

        private void CloseMarkerInfo()
        {
            if (CurrentlyDisplayedEndPointMarkerForm != null)
                CurrentlyDisplayedEndPointMarkerForm.Close();
            if (CurrentlyDisplayedRouteMarkerForm != null)
                CurrentlyDisplayedRouteMarkerForm.Close();
        }

        private void ClearMarkerInfo()
        {
            CurrentlyDisplayedEndPointMarkerForm = null;
            CurrentlyDisplayedRouteMarkerForm = null;
        }

        private void DeselectMarker()
        {
            if (CurrentlySelectedMarker == null) return;
            CurrentlySelectedMarker.IsSelected = false;
            CurrentlySelectedMarker = null;
        }

        private void SelectMarker(GMapSelectableMarker marker)
        {
            if (CurrentlySelectedMarker != null)
            {
                CurrentlySelectedMarker.IsSelected = false;
            }
            CurrentlySelectedMarker = marker;
            if (CurrentlySelectedMarker != null)
            {
                CurrentlySelectedMarker.IsSelected = true;
            }
            RedrawMap();
        }

        void MarkerClicked(GMapMarker mapMarker, MouseEventArgs e)
        {
            CloseMarkerInfo();
            ClearMarkerInfo();
            SelectMarker(mapMarker as GMapSelectableMarker);
            OpenMarkerInfo();
        }

        public void DoTraceRoute(IPAddress addr)
        {
            appList.ClearSelected();
            HideAllOverlays();

            GMapOverlay traceRTOverlay;
            if (ipOverlayTraceRTMap.ContainsKey(addr))
            {
                traceRTOverlay = ipOverlayTraceRTMap[addr].Item1;
                traceRTOverlay.IsVisibile = true;

                // Select item in trace route list
                tracesList.SelectedItem = ipOverlayTraceRTMap[addr].Item2;
            }
            else
            {
                traceRTOverlay = new GMapOverlay(mapexplr, "tracert" + addr);
                mapexplr.Overlays.Add(traceRTOverlay);

                // TODO: Maybe resolve host name from IP?
                StringListItem item = new StringListItem(addr.ToString(), new StringListItem.ListItemAction(delegate() {
                    appList.ClearSelected();
                    HideAllOverlays();
                    ipOverlayTraceRTMap[addr].Item1.IsVisibile = true;
                }));
                HostnameResolver hr = new HostnameResolver(addr);
                hr.HostnameResolvedEvent += new HostnameResolver.HostNameResolvedHandler(delegate(string hostname) {
                    item.Text = hostname + " (" + addr.ToString() + ")";
                });

                ipOverlayTraceRTMap.Add(addr, Tuple.Create(traceRTOverlay, item as IListItem));
                
                tracesList.Items.Add(item);
                tracesList.SelectedItem = item;

                int routersFound = 0;
                traceRTStatusLbl.Text = "Tracing route #" + routersFound;

                List<PointLatLng> routes = new List<PointLatLng>();

                Action<IPAddress, Location> addRouteToMapAction = new Action<IPAddress, Location>(
                    delegate(IPAddress route, Location loc) {
                        if (loc == null)
                            return;

                        PointLatLng point = new PointLatLng(loc.latitude, loc.longitude);
                        routes.Add(point);

                        // Update map
                        traceRTOverlay.Routes.Clear();
                        GMapRoute mapRoute = new GMapRoute(routes, "tracert" + addr);
                        mapRoute.Stroke = new Pen(Brushes.Red);
                        traceRTOverlay.Routes.Add(mapRoute);

                        traceRTOverlay.Markers.Add(new GMapRouteMarker(point, route));
                    }
                );

                TraceRoute tr = new TraceRoute(addr);
                tr.RouterFoundEvent += new TraceRoute.RouterFoundHandler(delegate(IPAddress route, Location loc) {
                    addRouteToMapAction(route, loc);
                    traceRTStatusLbl.Text = "Tracing route #" + (++routersFound);
                });
                tr.RouterFailedEvent += new TraceRoute.RouterFailedHandler(delegate() {
                    traceRTStatusLbl.Text = "Trace route to "+ (routersFound++) +" failed! Trying next.";
                });
                tr.TraceRouteCompletedEvent += new TraceRoute.TraceRouteCompletedHandler(delegate(IPAddress dest) {
                    if (dest == null || !dest.Equals(addr))
                    {
                        Location loc = GeoLookup.Get().LookupIPLocation(addr);
                        addRouteToMapAction(addr, loc);
                    }

                    traceRTStatusLbl.Text = "Trace route completed.";
                });
                tr.BeginTrace();
            }
        }

        private void appList_MeasureItem(object sender, MeasureItemEventArgs e)
        {
            e.ItemHeight = (int)(25f + 2 * 5);
        }

        private void DrawListItem(object sender, DrawItemEventArgs e)
        {
            if (e.Index == -1)
                return;
            
            // Get the ListBox and the item.
            ListBox lst = sender as ListBox;
            (lst.Items[e.Index] as IListItem).Draw(e);
        }

        public void ShowOverlayForProcess(Process process)
        {
            tracesList.ClearSelected();
            HideAllOverlays();

            if (process != null)
            {
                processToOverlayMap[process].IsVisibile = true;
            }
            else
            {
                unknownOverlay.IsVisibile = true;
            }
        }

        private void SelectedListItemChanged(object sender, EventArgs e)
        {
            ListBox lb = sender as ListBox;
            if (lb.SelectedItem != null)
                (lb.SelectedItem as IListItem).PerformAction(this);
        }

        private void MainWindow_FormClosing(object sender, FormClosingEventArgs e)
        {
            winPcap.StopCapture();
            worker.CancelAsync();
        }

        private void ArbitraryTraceRouteButton_Click(object sender, EventArgs e)
        {
            new ArbitraryTraceRouteForm(this).Show();
        }
    }
}
